from airflow.models.idata_workflow import WorkFlow


def gen_parallel_content(task_ids):
    data = {
        'nodes': []
    }
    start = {'x': 10, 'y': 10, 'type': 'START', 'stateId': 'start', 'name': 'START', 'branches': []}
    for tid in task_ids:
        task_node = {
            "x": 10,
            "y": 10,
            "type": "SHELL",
            "stateId": str(tid),
            "taskId": tid,
            "name": "task-[]".format(tid),
            "branches": [
                {
                    "branchId": "1",
                    "nextStateId": "end"
                }
            ]
        }
        data['nodes'].append(task_node)
        start['branches'].append({'branchId': '1', 'nextStateId': str(tid)})

    end = {'x': 10, 'y': 10, 'type': 'END', 'stateId': 'end', 'name': 'END', 'branches': []}
    data['nodes'].append(start)
    data['nodes'].append(end)
    return data


def gen_serial_content(task_ids):
    data = {
        'nodes': [{'x': 10, 'y': 10, 'type': 'START', 'stateId': 'start', 'name': 'START', 'branches': []}]
    }
    for tid in task_ids:
        task_node = {
            "x": 10,
            "y": 10,
            "type": "SHELL",
            "stateId": str(tid),
            "taskId": tid,
            "name": "task-{}".format(tid),
            "branches": []
        }
        data['nodes'][-1]['branches'].append({'branchId': '1', 'nextStateId': str(tid)})
        data['nodes'].append(task_node)
    data['nodes'][-1]['branches'].append({'branchId': '1', 'nextStateId': 'end'})
    end = {'x': 10, 'y': 10, 'type': 'END', 'stateId': 'end', 'name': 'END', 'branches': []}
    data['nodes'].append(end)
    return data


def update_content(workflow_id, content):
    workflow = WorkFlow.get_workflow(workflow_id)
    workflow.update_content(content)


if __name__ == '__main__':
    import json

    data = gen_serial_content([2, 3])
    data2 = {"nodes": [{"branches":[{"nextStateId":"17313282-0bb6-4d73-baf3-f237a85b5f15","branchId":1}],"name":"Start","nodeId":"1","stateId":"98ebc349-5240-4b25-acbb-fcc1b8267a8a","type":"START","x":610,"y":92},{"nodeId":1,"type":"SHELL","name":"Shell","x":360,"y":195,"stateId":"17313282-0bb6-4d73-baf3-f237a85b5f15","branches":[],"taskId":62}]}
    update_content(34, data2)
    # print(json.dumps(data, indent=4))
